#include "kernel_operator.h"
#define BUFFER_NUM 2
#define BLOCK_SIZE 256
using namespace AscendC;

template <typename T>
class KernelSelectV2 {
    public:
    __aicore__ inline KernelSelectV2() {}
    __aicore__ inline void Init(GM_ADDR condition, GM_ADDR x1,
                               GM_ADDR x2, GM_ADDR y,
                                uint32_t preiter, uint32_t totalBlocks, uint32_t totalLength, TPipe *pipeIn) {
	this->pipe = pipeIn;
        int loop = totalBlocks/preiter + (totalBlocks%preiter > 0);
        int coreDataNum = totalBlocks/preiter*(32/sizeof(T));
        int tailDataNum = coreDataNum;
        if (totalLength != totalBlocks*(32/sizeof(T))) {
            tailDataNum = totalLength - (loop - 1)*coreDataNum;
        }

        conditionGm.SetGlobalBuffer((__gm__ int8_t *)condition, totalLength);
        x1Gm.SetGlobalBuffer((__gm__ T *)x1, totalLength);
        x2Gm.SetGlobalBuffer((__gm__ T *)x2, totalLength);
        yGm.SetGlobalBuffer((__gm__ T *)y, totalLength);
        this->elemCount = (256/sizeof(T));

        pipe->InitBuffer(conditionQueue, BUFFER_NUM, coreDataNum*sizeof(int8_t));
        pipe->InitBuffer(x1Queue, BUFFER_NUM, coreDataNum*sizeof(T));
        pipe->InitBuffer(x2Queue, BUFFER_NUM, coreDataNum*sizeof(T));
        pipe->InitBuffer(yQueue, BUFFER_NUM, coreDataNum*sizeof(T));
        pipe->InitBuffer(maskbuf, (coreDataNum+255)/256*256*sizeof(int8_t));
        pipe->InitBuffer(conditionCast,(coreDataNum+elemCount-1)/elemCount*elemCount*sizeof(int32_t));

    }
    __aicore__ inline void process() {
        for (int i = 0; i < loop - 1; i++) {
            copyIn(i, coreDataNum);
           compute(i, coreDataNum);
           copyOut(i, coreDataNum);
        }
        copyIn(loop-1, coreDataNum);
        compute(loop-1, coreDataNum);
        copyOut(loop-1, coreDataNum);
    }
    __aicore__ inline void copyIn(int i, int count) {
        auto x1Local = x1Queue.AllocTensor<T>();
        auto x2Local = x2Queue.AllocTensor<T>();
        auto conditionLocal = conditionQueue.AllocTensor<int8_t>();
        AscendC::DataCopyExtParams copyParams{1, static_cast<uint32_t>(count*sizeof(T)), 0, 0, 0}; // 结构体DataCopyExtParams最后一个参数是rsv保留位
        AscendC::DataCopyPadExtParams<T> padParams{false, 0, 0, 0};
        DataCopyPad(x1Local, x1Gm[i*coreDataNum], copyParams, padParams);
        DataCopyPad(x2Local, x2Gm[i*coreDataNum], copyParams, padParams);
        AscendC::DataCopyExtParams copyParamsBool{1,static_cast<uint32_t>(count), 0, 0, 0}; // 结构体DataCopyExtParams最后一个参数是rsv保留位
        AscendC::DataCopyPadExtParams<int8_t> padParamsBool{false, 0, 0, 0};
        AscendC::DataCopyPad(conditionLocal, conditionGm[i*coreDataNum], copyParamsBool, padParamsBool); // 从GM->VECIN搬运40Bytes

        x1Queue.EnQue(x1Local);
        x2Queue.EnQue(x2Local);
        conditionQueue.EnQue(conditionLocal);
    }

    __aicore__ inline void compute(int i, int count) {
        int roundCount = (roundCount + elemCount-1)/elemCount*elemCount;
        auto x1Local = x1Queue.DeQue<T>();
        auto x2Local = x2Queue.DeQue<T>();
        auto yLocal = yQueue.DeQue<T>();
        auto conditionLocal = conditionQueue.DeQue<int8_t>();
        auto maskLocal = maskbuf.Get<T>();
        auto conditionCastLocal = conditionCast.Get<int16_t>();
        AscendC::Cast(conditionCastLocal, conditionLocal, AscendC::RoundMode::CAST_CEIL, roundCount);
	    CompareScalar(maskLocal, conditionCastLocal, (int16_t)1, AscendC::CMPMODE::EQ, roundCount);
        AscendC::Select(yLocal, maskLocal, x1Local, x2Local, SELMODE::VSEL_TENSOR_TENSOR_MODE, roundCount);

        yQueue.EnQue(yLocal);
        x1Queue.FreeTensor(x1Local);
        x2Queue.FreeTensor(x2Local);
        yQueue.FreeTensor(yLocal);

        conditionQueue.FreeTensor(conditionLocal);
    }

    __aicore__ inline void copyOut(int i, int count) {
        auto yLocal = yQueue.DeQue<T>();
        AscendC::DataCopyExtParams copyParams{1, static_cast<uint32_t>(count * sizeof(T)), 0, 0, 0};
        DataCopyPad(yGm[i*coreDataNum],yLocal, copyParams);
    }


private:
    TPipe *pipe;
    GlobalTensor<T> yGm, x1Gm, x2Gm;
    GlobalTensor<int8_t> conditionGm;
    TQue<QuePosition::VECIN, BUFFER_NUM> conditionQueue;
    TQue<QuePosition::VECIN, BUFFER_NUM> x1Queue,x2Queue;
    TQue<QuePosition::VECOUT, BUFFER_NUM> yQueue;
    TBuf<QuePosition::VECCALC> castx1Buf, castx2Buf, castyBuf, maskbuf, conditionCast;
    int loop,  coreDataNum,tailDataNum, elemCount;
};

template <typename T>
class KernelSelectCastV2 {
    public:
    __aicore__ inline KernelSelectCastV2() {}
    __aicore__ inline void Init(GM_ADDR condition, GM_ADDR x1,
                               GM_ADDR x2, GM_ADDR y,
                                uint32_t preiter, uint32_t totalBlocks, uint32_t totalLength, TPipe* pipeIn) {
	this->pipe = pipeIn;
        int loop = totalBlocks/preiter + (totalBlocks%preiter > 0);
        int coreDataNum = totalBlocks/preiter*(32/sizeof(T));
        int tailDataNum = coreDataNum;
        if (totalLength != totalBlocks*(32/sizeof(T))) {
            tailDataNum = totalLength - (loop - 1)*coreDataNum;
        }
        this->elemCount = (256/sizeof(T));

        conditionGm.SetGlobalBuffer((__gm__ int8_t *)condition, totalLength);
        x1Gm.SetGlobalBuffer((__gm__ T *)x1, totalLength);
        x2Gm.SetGlobalBuffer((__gm__ T *)x2, totalLength);
        yGm.SetGlobalBuffer((__gm__ T *)y, totalLength);

        pipe->InitBuffer(conditionQueue, BUFFER_NUM, coreDataNum*sizeof(bool));
        pipe->InitBuffer(x1Queue, BUFFER_NUM, coreDataNum*sizeof(T));
        pipe->InitBuffer(x2Queue, BUFFER_NUM, coreDataNum*sizeof(T));
        pipe->InitBuffer(yQueue, BUFFER_NUM, coreDataNum*sizeof(T));
        pipe->InitBuffer(maskbuf, (coreDataNum+7)/8*sizeof(int8_t));

        if constexpr(std::is_same_v<DTYPE_X1, int8_t>) {
            pipe->InitBuffer(cast1buf, coreDataNum*sizeof(half));
            pipe->InitBuffer(cast2buf, coreDataNum*sizeof(half));
            pipe->InitBuffer(cast3buf, coreDataNum*sizeof(half));

        }
        if constexpr(std::is_same_v<DTYPE_X1, int32_t>) {
            pipe->InitBuffer(cast1buf, coreDataNum*sizeof(float));
            pipe->InitBuffer(cast2buf, coreDataNum*sizeof(float));
            pipe->InitBuffer(cast3buf, coreDataNum*sizeof(float));
        }
    }
    __aicore__ inline void process() {
        for (int i = 0; i < loop - 1; i++) {
            copyIn(i, coreDataNum);
            compute(i, coreDataNum);
            copyOut(i, coreDataNum);
        }
        copyIn(loop-1, coreDataNum);
        compute(loop-1, coreDataNum);
        copyOut(loop-1, coreDataNum);
    }
    __aicore__ inline void copyIn(int i, int count) {
        auto x1Local = x1Queue.AllocTensor<T>();
        auto x2Local = x2Queue.AllocTensor<T>();
        auto conditionLocal = conditionQueue.AllocTensor<bool>();
        AscendC::DataCopyExtParams copyParams{1, static_cast<uint32_t>(count*sizeof(T)), 0, 0, 0}; // 结构体DataCopyExtParams最后一个参数是rsv保留位
        AscendC::DataCopyPadExtParams<T> padParams{false, 0, 0, 0};
    	AscendC::LocalTensor<int8_t> interpreTensor = conditionLocal.ReinterpretCast<int8_t>();

        DatacopyPad(x1Local, x1Gm[i*coreDataNum], copyParams, padParams);
        DatacopyPad(x2Local, x2Gm[i*coreDataNum], copyParams, padParams);
        AscendC::DataCopyExtParams copyParamsBool{1, static_cast<uint32_t>(count), 0, 0, 0}; // 结构体DataCopyExtParams最后一个参数是rsv保留位
        AscendC::DataCopyPadExtParams<int8_t> padParamsBool{false, 0, 0, 0};
        AscendC::DataCopyPad(interpreTensor, conditionGm[i*coreDataNum], copyParamsBool, padParamsBool); // 从GM->VECIN搬运40Bytes
        x1Queue.EnQue(x1Local);
        x2Queue.EnQue(x2Local);
        conditionQueue.EnQue(conditionLocal);
    }

    __aicore__ inline void compute(int i, int count) {
        auto x1Local = x1Queue.DeQue<T>();
        auto x2Local = x2Queue.DeQue<T>();
        auto yLocal = yQueue.DeQue<T>();
        if ()
        auto cast1Local = cast1buf.Get<T>();
        auto cast2Local = cast2buf.Get<T>();
        auto cast3Local = cast3buf.Get<T>();
       AscendC::Cast(cast1Local, x1Local, AscendC::RoundMode::CAST_CEIL, count);
       AscendC::Cast(cast2Local, x2Local, AscendC::RoundMode::CAST_CEIL, count);
       AscendC::Cast(cast3Local, yLocal, AscendC::RoundMode::CAST_CEIL, count);

        auto conditionLocal = conditionQueue.DeQue<bool>();
        auto maskLocal = maskbuf.Get<T>();
       CompareScalar(maskLocal, conditionLocal, (bool)1, AscendC::CMPMODE::EQ, count);
       AscendC::Select(cast3Local, maskLocal, cast1Local, cast2Local, SELMODE::VSEL_TENSOR_TENSOR_MODE, count);
       AscendC::Cast(yLocal, cast3Local, AscendC::RoundMode::CAST_CEIL, count);

        yQueue.EnQue(yLocal);
        x1Queue.FreeTensor(x1Local);
        x2Queue.FreeTensor(x2Local);
        yQueue.FreeTensor(yLocal);

        conditionQueue.FreeTensor(conditionLocal);
    }

    __aicore__ inline void copyOut(int i, int count) {
        auto yLocal = yQueue.DeQue<T>();

        AscendC::DataCopyExtParams copyParams{1, static_cast<uint32_t>(count * sizeof(T)), 0, 0, 0};
        DataCopyPad(yGm[i*coreDataNum],yLocal, copyParams);
    }


private:
    TPipe *pipe;
    GlobalTensor<T> yGm, x1Gm, x2Gm;
    GlobalTensor<int8_t> conditionGm;
    TQue<QuePosition::VECIN, BUFFER_NUM> conditionQueue;
    TQue<QuePosition::VECIN, BUFFER_NUM> x1Queue,x2Queue;
    TQue<QuePosition::VECOUT, BUFFER_NUM> yQueue;
    TBuf<QuePosition::VECCALC> cast1buf, cast2buf, cast3buf, maskbuf;
    int loop,  coreDataNum,tailDataNum;

};
extern "C" __global__ __aicore__ void select_v2(GM_ADDR condition, GM_ADDR x1, GM_ADDR x2, GM_ADDR y, GM_ADDR workspace, GM_ADDR tiling) {
    GET_TILING_DATA(tiling_data, tiling);
      TPipe pipe;

    // TODO: user kernel impl
    if (TILING_KEY_IS(1)) {
        KernelSelectV2<DTYPE_X1>op;
        op.Init(condition, x1, x2,  y, tiling_data.preiter, tiling_data.totalBlocks, tiling_data.totalLength, &pipe);
        op.process();
    }
    if (TILING_KEY_IS(2)) {
//        KernelSelectCastV2<DTYPE_X1>op;
//        op.Init(condition, x1, x2,  y, tiling_data.preiter, tiling_data.totalBlocks, tiling_data.totalLength, &pipe);
//        op.process();
    }
}

